import pandas as pd
import numpy as np

# 读取 CSV 文件
df = pd.read_csv("data/new_file.csv")

# 提取 fhjl_time 列的前四位数字并转换为字符串
df['fhjl_time'] = df['fhjl_time'].astype(str)
df['extracted_time'] = df['fhjl_time'].str[:4]

# 计算相同提取时间的 hk 的和
hk_sum_by_time = df.groupby('extracted_time')['hk'].sum().reset_index()

# 将 extracted_time 列转换为整数类型
hk_sum_by_time['extracted_time'] = hk_sum_by_time['extracted_time'].astype(int)

# 计算增长率
hk_sum_by_time['growth_rate'] = hk_sum_by_time['hk'].pct_change()

# 计算平均增长率
average_growth_rate = hk_sum_by_time['growth_rate'].mean()

# 预测未来的增长趋势
future_times = np.arange(int(hk_sum_by_time['extracted_time'].max()) + 1, int(hk_sum_by_time['extracted_time'].max()) + 6)
hk_sum_by_time['extracted_time'] = hk_sum_by_time['extracted_time'].astype(int)  # 再次确保 extracted_time 为整数类型
predicted_hk = hk_sum_by_time['hk'].iloc[-1] * (1 + average_growth_rate) ** (future_times - hk_sum_by_time['extracted_time'].iloc[-1])

# 打印结果
print("平均增长率：", average_growth_rate)
print("未来 hk 的预测值：")
for time, hk in zip(future_times, predicted_hk):
    print(f"时间 {time}: {hk}")

# 绘制图形
import matplotlib.pyplot as plt

plt.plot(hk_sum_by_time['extracted_time'], hk_sum_by_time['hk'], label='实际数据')
plt.plot(future_times, predicted_hk, label='预测趋势')
plt.xlabel('提取的时间前四位')
plt.ylabel('hk 的和')
plt.title('hk 的和的增长趋势预测')
plt.legend()
plt.show()
